function out=momentszselstate(GG,CC,RR,SDX,ZZ,in,pos)
% =========================================================================
% function
% out=momentszselstate(GG,CC,RR,SDX,ZZ,case_tag,in,pos)
%
% MOMENTZSELSTATE.M 
%
%% 1. Purpose. This is a generalization of MODE_MOMENTSALL 
% As that function, it computes moments for all the observables in 
% Z, as well as rows of the state vector given
% in POS, a position indicator. 
% 
%  
%% 2. Inputs in IN structure 
%
%  Assign default values ( ) if not provided 
%
%  ========================================================================
%
% IN.VECSIM:       [NDROP NTAKE] ([0 0]) 
%
% IN.NREP:         Number of replications (100) 
%
% IN.NPER:         Number of periods for the Autocorrelation/Cross
%                  Correlation (4) 
%
% IN.NIRF:         Number of horizons for IRFS (16) 
%
% IN.NVDEC:         Number of horizons for variance decompositions (40) 
%
% IN.FLAG_UNIT:     Binary, equal 1 if report unit impulse, 0 if 1 std impulse (0)
%
% IN.FLAG_COV:       Binary, equal 1 if covariances, o if correlations (0).
% Applies only to simulated moments. Both theoretical moments will be
% computed. 
%
% IN.FLAG_ADDZ:      [NZ 1] Binary vector, i-th row 1 if add to form a
% level  (no default, must be defined) 
%
% IN.FLAG_ADDS:      [NS 1] Binary cector, i-th row 1 if add to form a level 
%                   where NS is the dimension in POS, not of the full state
%                   space (no default, must be defined)
%
% IN.FLAG_MEDBONLY:  Binary, equal 1, will only report medians and [5 95] bands (1)
% 
%.IN.FLAG_ACDECOMP: Binary, equal 1, will decompose AC for observables (0);
%
%% 3. Shutting down some computations
% 
% Set: 
%
% IN.NIRF == 0    No IRF
% 
% IN.NVDEC== 0    Only VDECSTA 
% 
% IN.VECSIM(2)=0  No simulated moments 
% 
% IN.NPER ==0     No AC or COV, simulated or Asymptotic 
%
%% 4. OUT structure
% 
% Ending and NR (number of rows) varies: 
% 
% 1) Ending Z applies to the observable Z matrix, NR=NZ 
%
% 2) Ending S applies to the selected states, NR=NS. If no states requested
% then [] matrix will be passed. 
%
% I.) Size Theoretical (Population) Moments 
% 
% =========================================================================
%
%
% .VCVTH:     [NR NR] theoretical Variance Covariance Matrix 
%
% .STDTH:     [NT  1] theoretical Std vector for observables
%
% .ACTH:      [NR  NR NPER+1] theoretical Autocorrelation for observables
%
% .COVTH:     [NR  NR NPER+1] theoretical Autocovariances for observables
%
% .IRF:      [NR  NR NIRF+1] impulse responses
%
% .VDECSTA:   [NR NX]   Variance decomposition, stationary
%                            (asymptotic)
% .VDECH:      [NR NX NVDEC]   Variance decomposition, horizon 1 through 
%                             NVDEC
%
% .ACTHZ_DECOMP: [NZ NZ NPER+1 NX} Decomposition of the correlation matrix 
%               across all shocks. ONLY FOR the observables 
%
% II) Size Simulated (Sample) Moments 
% 
% =========================================================================
%
% NPAGE dimension depends on FLAG_MEDBONLY: 
%
%   == 0    NREP  All replications 
%
%   == 1    3  Median and [5,95] bands 
%
% .STDSIM     [NR NPAGE]
%
% .ACSIM      [NR NR NPAGE+1] if FLAG_ACOV=1
% 
% .COVSIM     [NR NR NPAGE+1] if FLAG_ACOV=0 
% 
% Alejandro Justiniano (C)
%
% Created Oct 2 2009 
%
% Modified: June 29 2010 
% =========================================================================

%% 1. Check dimension and default values 

% I. Dimensions, part 1 
% =========================================================================
if nargin < 7; error('Need 7 inputs');end 
ny=size(GG,1);
nx=size(SDX,1);
nz=size(ZZ,1);
if isempty(pos)
    ns=0;
else
    ns=length(pos);
end

% II. Default values 
% =========================================================================
[junk,in]=ch_field(in,'vecsim'  ,[0 0]); 
[junk,in]=ch_field(in,'nrep'    ,500   ); 
[junk,in]=ch_field(in,'nper'    ,4   ); 
[junk,in]=ch_field(in,'nirf'    ,16  );
[junk,in]=ch_field(in,'nvdec'   ,40  );
[junk,in]=ch_field(in,'flag_unit',0  );
[junk,in]=ch_field(in,'flag_acov' ,0 ); 
[junk,in]=ch_field(in,'flag_medbonly',1); 
[junk,in]=ch_field(in,'flag_acdecomp',0); 

% III. Dimensions, part 2 
% =========================================================================
if in.nirf > 0
    if length(in.flag_addz)~=nz
        error('in.FLAG_ADDZ must be NZ x 1')
    end
    if ns > 0
        if length(in.flag_adds)~=ns
            error('in.FLAG_ADDS must be NS x 1')
        end
    end
end 

if in.flag_medbonly~=0 && in.flag_medbonly~=1 
    error('in.flag_medbonly must be 0 or 1') 
end 
if in.flag_acov~=0 && in.flag_acov~=1;
    error('FLAG_ACOV must be 0 or 1'); 
end 
if in.flag_acov~=0 && in.flag_acov~=1;
    error('FLAG_ACOV must be 0 or 1'); 
end 


%% 2. Theoretical (Monthly) Population Moments

% 1A. Variance, Covariance and Autocorrelation at horizon 0 through NPER
% =========================================================================

% I. VCVTH [NR NR] Theoretical Variances 
% =========================================================================
pshat=disclyap_fast(GG,RR*(SDX'*SDX)*RR');
if ns > 0
    vcvths=pshat(pos,pos);
else
    vcvths=[];
end
vcvthz=ZZ*pshat*ZZ';

% II. STDTH [NR  1] Theoretical Standard Deviations 
% =========================================================================
out.stdthz=sqrt(diag(vcvthz));
if ns > 0
    out.stdths=sqrt(diag(vcvths));
    denos=out.stdths*out.stdths'; 
else
    out.stdths=[];
end
deno=out.stdthz*out.stdthz';

% III. ACTH [NR 1] Theoretical Correlations 
% =========================================================================
out.acthz= zeros(nz,nz,in.nper+1);
if ns > 9
    out.acths= zeros(ns,ns,in.nper+1);
else
    out.acths=[];
end

% IV. COVTH [NR  NR NPER+1] Theoretical Covariances 
% =========================================================================
out.covthz=zeros(nz,nz,in.nper+1);
out.covthz(:,:,1)=vcvthz;
if ns > 0
    out.covths=zeros(ns,ns,in.nper+1);
    out.covths(:,:,1)=vcvths; 
else
    out.covths=[];
end

% V. VDECSTA [NR NX] Stationarty Variance decomposition
%    Declaration 
% =========================================================================
out.vdecstaz=zeros(nz,nx);
if ns > 0 
out.vdecstas=zeros(ns,nx); 
else 
    out.vdecstas=[]; 
end 
pall=zeros(ny,ny,nx); 

% VI. ACTH [NR  NR NPER+1] Theoretical Correlation 
% At (h=0) 
% =========================================================================
out.acthz(:,:,1)     =squeeze(out.covthz(:,:,1))./deno;
if ns > 0
    out.acths(:,:,1) =squeeze(out.covths(:,:,1))./denos;
end

% VII. ACTHZ_DECOMP [NZ NZ NPER+1 NX} Decomposition of the correlation matrix 
% across all shocks. Note: ONLY FOR the observables 
% =========================================================================
if in.flag_acdecomp==1;
    out.acthz_decomp=zeros(nz,nz,in.nper+1,nx);
end

% VDECSTA [NR NX] Stationarty Variance decomposition, Computation 
% -------------------------------------------------------------------------
for ii=1:nx;
    
    pall(:,:,ii)=real(disclyap_fast(GG,RR*((SDX(ii,:)' )*SDX(ii,:))*RR'));

    out.vdecstaz(:,ii)=diag(squeeze(ZZ*pall(:,:,ii)*ZZ'))./(diag(vcvthz));    
    if ns > 0
        out.vdecstas(:,ii)=diag( pall(pos,pos,ii) )./(diag(vcvths));
    end
        
    if in.flag_acdecomp==1;
        out.acthz_decomp(:,:,1,ii)=squeeze(ZZ*pall(:,:,ii)*ZZ')./deno; 
    end 
    
end
    
% COVTH and ACTH Covariances and Correlations at h=1,....,NPER horizon 
% -------------------------------------------------------------------------
if in.nper > 0
    Gprod=eye(ny);
    for ii=1:in.nper;
        Gprod=Gprod*GG;
        temp=Gprod*pshat;        
        % Covariances and Correlations 
        if ns > 0
            out.covths(:,:,ii+1)=temp(pos,pos);
            out.acths(:,:,ii+1 )=squeeze(out.covths(:,:,ii+1))./denos;
        end
        
        out.covthz(:,:,ii+1)=ZZ*temp*ZZ';
        out.acthz(:,:,ii+1 )=squeeze(out.covthz(:,:,ii+1))./deno;
        
        if in.flag_acdecomp==1
            for mm=1:nx;
                out.acthz_decomp(:,:,ii+1,mm)=(ZZ*Gprod*squeeze(pall(:,:,mm))*(ZZ'))./(deno);
            end
        end
    end
    
    % Check ACTH and ACTH_DECOMP add-up 
    if in.flag_acdecomp==1        
        if max(max(max(abs(out.acthz-sum(out.acthz_decomp,4))))) > 1e-8;
            error('ACTH_DECOMP does not recover the total AC')
        end
    end
end
clear temp; 

if in.flag_acdecomp==1;
    out.acth_decomp=acth_decomp;
end
    
if in.flag_unit~=1 && in.flag_unit~=0
    in.flag_unit=0;
end

% VII. Impulse responses at horizon 0 through NIRF
% =========================================================================
if in.nirf~=0
    
    indic_addz=find(in.flag_addz==1);
    out.irfz=zeros(nz,nx,in.nirf+1);

    if ns > 0
        indic_adds=find(in.flag_adds==1);
        out.irfs=zeros(ns,nx,in.nirf+1);
    else 
        out.irfs=[]; 
    end 
    
    for ii=1:nx
        
        if in.flag_unit==0
            resp=SDX(ii,ii)*RR(:,ii);
        else
            resp=RR(:,ii);
        end
        if ns > 0
            out.irfs(:,ii,1)=resp(pos);
        end
        
        out.irfz(:,ii,1)=ZZ*resp;
        for mm=2:(in.nirf+1);
            resp=GG*resp;
            if ns > 0 
            out.irfs(:,ii,mm)=resp(pos);
            end 
            out.irfz(:,ii,mm)=ZZ*resp;
        end
        
        if ~isempty(indic_addz);
            out.irfz(indic_addz,ii,:)=cumsum(out.irfz(indic_addz,ii,:),3);
        end
        if ns > 0
            if ~isempty(indic_adds);
                out.irfs(indic_adds,ii,:)=cumsum(out.irfs(indic_adds,ii,:),3);
            end
        end
        
    end
    
else 
    out.irfz=[]; 
    out.irfs=[]; 
end 
    
% VII. Variance decompositions at horizon 0 through NVDEC
% =========================================================================
if in.nvdec==0;
    out.vdechz=[];
    out.vdechs=[]; 
else
    % Variance decomposition at h=1,....,NPER horizon 
    [vdhhz,vdhhs]=mode_momentsall_vdec(in.nvdec);
    out.vdechz=vdhhz;
    %out.vdechs=zeros(ns,nx,in.nvdec);
    if ns > 0
        out.vdechs=vdhhs(pos,:,:);
    end
end

% =========================================================================

%% 3. Simulated (Monthly) Sample Moments, ALL 

% Recall that if FLAG_MEDBONLY==1 it will only report the median and 
% [5 95] percentile bands of all of the simulated moments 

% =========================================================================

if in.vecsim(2)==0 
     out.stdsimz =[]; 
     out.stdsims =[]; 
     out.acsimz  =[];
     out.acsims  =[]; 
     out.covsimz =[]; 
     out.covsims =[]; 
     return
end

nobs=in.vecsim(1)+in.vecsim(2);

out.stdsimz=zeros(nz,in.nrep);
if ns > 0
    out.stdsims=zeros(ns,in.nrep);
end

% Storage for simulated COV or CORR only 
% =========================================================================
if in.nper > 0
    if in.flag_acov==0;
        
        out.acsimz=zeros(nz,nz,in.nper+1,in.nrep);
        if ns > 0 
        out.acsims=zeros(ns,ns,in.nper+1,in.nrep); 
        else 
            out.acsims=[]; 
        end 
        out.covsimz=[];
        out.covsims=[]; 
        
    else
        
        out.covsimz=zeros(nz,nz,in.nper+1,in.nrep);
        if ns > 0
            out.covsims=zeros(ns,ns,in.nper+1,in.nrep);
        else
            out.covsims=[];
        end
        out.acsimz=[];
        out.acsims=[]; 
        
    end
end

% Beging Simulation 
% =========================================================================
for hh=1:in.nrep;
    
    % Simulate shocks, states and observable data 
    % --------------------------------------------
    shocks=(randn(nobs,nx)*SDX)';
    ssim=zeros(size(GG,1),nobs);
    ysim=zeros(size(ZZ,1),nobs);
    ssim(:,1)=RR*shocks(:,1);
    ysim(:,1)=ZZ*(ssim(:,1) + CC);
    for ii=2:nobs;
        ssim(:,ii)=GG*ssim(:,ii-1)+RR*shocks(:,ii);
        ysim(:,ii)=ZZ*(ssim(:,ii) + CC);
    end
    
    ysim=ysim(:,in.vecsim(1)+1:end)';
    ssim=ssim(:,in.vecsim(1)+1:end)';
    
    out.stdsimz(:,hh)=(std(ysim)');
    if ns > 0 
    out.stdsims(:,hh)=(std(ssim(:,pos))');
    end 
    
    % Compute covariances OR correlations 
    % ACCOV has two outputs 
    % 1. [nz nz nper+1] matrix of correlations 
    % 2. [nz nz nper+1] matrix of covariances  
    % =====================================================================
    if in.nper > 0
        if in.flag_acov==1
            [junk,out.covsimz(:,:,:,hh)]=accov(ysim,(0:in.nper));
            if ns > 0 
            [junk,out.covsims(:,:,:,hh)]=accov(ssim(:,pos),(0:in.nper));
            end 
        else
            [out.acsimz(:,:,:,hh)]=accov(ysim,(0:in.nper));
            if ns > 0 
            [out.acsims(:,:,:,hh)]=accov(ssim(:,pos),(0:in.nper));
            end 
        end
    end                
        
end

%% 4. Simulated (Monthly) Sample Moments, Medians and Bands ONLY
% IF FLAG_MEDBONLY ==1 report only median and 5-95 bands
if in.flag_medbonly==1
    
    pvec=round(in.nrep*[0.05 0.5 0.95]);
    
    out.stdsimz=sort(out.stdsimz,2);
    out.stdsimz=out.stdsimz(:,pvec);
    
    if ns > 0
        out.stdsims=sort(out.stdsims,2);
        out.stdsims=out.stdsims(:,pvec);
    end
    
    if in.flag_acov==1
        out.covsimz=sort(out.covsimz,4);
        out.covsimz=out.covsimz(:,:,:,pvec);
        
        if ns > 0
            out.covsims=sort(out.covsims,4);
            out.covsims=out.covsims(:,:,:,pvec);
        end
    else
        out.acsimz=sort(out.acsimz,4);
        out.acsimz=out.acsimz(:,:,:,pvec);
        if ns > 0
            out.acsims=sort(out.acsims,4);
            out.acsims=out.acsims(:,:,:,pvec);
        end
    end
    
end

% =========================================================================
% Variance covariance decomposition across horizons 
% =========================================================================
    function [vdech,vdech_st]=mode_momentsall_vdec(nvdec) 
        % VDECH_ST is the decomposition for the WHOLE state vector 
        
        Gnew=eye(ny);
        
        vdech=zeros(nz,nx,nvdec);
        vdech_st=zeros(ny,nx,nvdec);
        
        vzero=RR*(SDX')*SDX*(RR');
        vcum=zeros(nz,nz,nvdec,nx);
        vcum_st=zeros(ny,ny,nvdec,nx);
        for kk=1:nvdec;
            % ====================
            % Whole Variance
            % ====================
            if kk==1;
                vcfull_st=vzero;
                vcfull=ZZ*vcfull_st*(ZZ');
            else
                Gnew=GG*Gnew ;
                vcfull_st=Gnew*vzero*(Gnew')+vcfull_st;
                vcfull=ZZ*Gnew*vzero*(Gnew')*(ZZ')+vcfull;
            end;
            deno_temp=diag(vcfull);
            deno_st=diag(vcfull_st);
            if kk==1;
                for jj=1:nx
                    et=( SDX(jj,:)' )*SDX(jj,:);
                    vcum_st(:,:,1,jj)=(RR*et*(RR'));
                    vcum(:,:,1,jj)=ZZ*((RR*et*(RR')))*(ZZ');
                    vdech(:,jj,kk)=(diag(vcum(:,:,1,jj)))./deno_temp;
                    vdech_st(:,jj,kk)=(diag(vcum_st(:,:,1,jj)))./deno_st;
                end;
            else
                for jj=1:nx;
                    et=( SDX(jj,:)' )*SDX(jj,:);
                    vcum(:,:,kk,jj)=ZZ*(Gnew*RR*et*(RR')*(Gnew'))*(ZZ')+vcum(:,:,kk-1,jj);
                    vcum_st(:,:,kk,jj)=(Gnew*RR*et*(RR')*(Gnew'))+vcum_st(:,:,kk-1,jj);
                    vdech(:,jj,kk)=(diag(vcum(:,:,kk,jj)))./deno_temp;
                    vdech_st(:,jj,kk)=(diag(vcum_st(:,:,kk,jj)))./deno_st;
                end
            end
            if any( sum( squeeze(vdech(:,:,kk)) , 2 ) < 0.99 )
                error('Variance Decompositions Observables are innacurate');
            end
            
            if any( sum( squeeze(vdech_st(:,:,kk)) , 2 ) < 0.99 )
                error('Variance Decompositions Observables are innacurate');
            end

            
        end       
    end

end 